from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough

from AliUtils import AliEmbeddings, get_model

model = get_model()
# 准备测试数据 ，假设我们提供的文档数据如下：
documents = [
	Document(
		page_content="狗是伟大的伴侣，以其忠诚和友好而闻名。",
		metadata={"source": "哺乳动物宠物文档"},
	),
	Document(
		page_content="猫是独立的宠物，通常喜欢自己的空间。",
		metadata={"source": "哺乳动物宠物文档"},
	),
	Document(
		page_content="金鱼是初学者的流行宠物，需要相对简单的护理。",
		metadata={"source": "鱼类宠物文档"},
	),
	Document(
		page_content="鹦鹉是聪明的鸟类，能够模仿人类的语言。",
		metadata={"source": "鸟类宠物文档"},
	),
	Document(
		page_content="兔子是社交动物，需要足够的空间跳跃。",
		metadata={"source": "哺乳动物宠物文档"},
	),
]

# 使用自定义Embeddings
vector_store = Chroma.from_documents(
	documents,
	embedding=AliEmbeddings()  # 替换为自定义类
)

# 检索器: bind(k=1) 返回相似度最高的第一个
retriever = RunnableLambda(vector_store.similarity_search).bind(k=1)

# print(retriever.batch(['咖啡猫', '鲨鱼']))

# 提示模板
message = """
使用提供的上下文仅回答这个问题:
{question}
上下文:
{context}
"""

prompt_template = ChatPromptTemplate([('human', message)])

# RunnablePassthrough允许我们将用户的问题之后再传递给prompt和model。
chain = {'question': RunnablePassthrough(), 'context': retriever} | prompt_template | model

print(chain.invoke('请介绍一下猫 ！').content)
